package com.archgeek.bigdata.presto;

import com.facebook.presto.jdbc.PrestoResultSet;
import com.facebook.presto.jdbc.PrestoStatement;
import com.facebook.presto.jdbc.QueryStats;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

import static com.google.common.base.Preconditions.checkState;

/**
 * @author pizhihui
 * @date 2022-03-30 15:02
 */
public class PrestoJdbcTest {

    private static ListeningExecutorService pool = MoreExecutors.listeningDecorator(new ThreadPoolExecutor(2, 10, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<Runnable>(1024),
            new ThreadFactoryBuilder()
                    .setNameFormat("presto-pool-%d").build(),
            new ThreadPoolExecutor.AbortPolicy()));

    public static void main(String[] args) {


        try {

//            String sql = "select f_p_date,count(1) from rt_mart.rt_app_reported_db_t_back_event group by f_p_date order by f_p_date desc";

            String sql = "create table dp_tmp.zhihuipi_20220616_test_01 stored as orc as \n" +
                    "select f_p_date,\n" +
                    "       count(1) as fnum\n" +
                    "from dp_ods.log_db_t_data_platform_log\n" +
                    "where f_p_date > '2022-01-10'\n" +
                    "group by f_p_date\n" +
                    "limit 10";

            Properties properties = new Properties();
            properties.put("user", "biadmin");
            properties.put("socksProxy", "127.0.0.1:1081");
            // Connection conn = DriverManager.getConnection("jdbc:presto://10.10.28.137:9000/hive?useProxy=true&proxyProtocol=socks5&proxyHost=127.0.0.1&proxyPort=1081", properties);
            Connection conn = DriverManager.getConnection("jdbc:presto://10.10.28.137:9000/hive", properties);
            Statement stmt = conn.createStatement();

            try {

                //stmt.setMaxRows(10);

                // ResultSet rs = stmt.executeQuery("explain " + sql);
                System.out.println("开始执行......");
                long start = System.currentTimeMillis();

                PrestoStatement prestoStatement = stmt.unwrap(PrestoStatement.class);
                RecordingProgressMonitor progressMonitor = new RecordingProgressMonitor();
                prestoStatement.setProgressMonitor(progressMonitor);

//                pool.execute(() -> {
//                    try {
//                        PrestoStatement prestoStatement = stmt.unwrap(PrestoStatement.class);
//                        RecordingProgressMonitor progressMonitor = new RecordingProgressMonitor();
//                        prestoStatement.setProgressMonitor(progressMonitor);
//                        prestoStatement.clearProgressMonitor();
//                        List<QueryStats> queryStatsList = progressMonitor.finish();
//                        if (!queryStatsList.isEmpty()) {
//                            String queryId = queryStatsList.get(0).getQueryId();
//                            System.out.println("执行 id: " + queryId);
//                        }
//
////                        prestoIdMap.put(idInt, queryId);
////                        DBUtills.updatePrestoUrl(idInt, String.format("http://10.10.28.137:9000/ui/query.html?%s", queryId));
//                    } catch (Exception e) {
//                        e.printStackTrace();
//                    }
//                });

                // execute
                boolean isRs = false;
                try {
                    isRs = stmt.execute(sql);
                    System.out.println("执行结束......." + isRs + ",时间: " + (System.currentTimeMillis() - start) / 1000);
                } catch (Exception e) {
                    throw e;
                } finally {
                    prestoStatement.clearProgressMonitor();
                    List<QueryStats> queryStatsList = progressMonitor.finish();
                    System.out.println("finall stats: " + queryStatsList.get(0).getQueryId());
                }



                ResultSet rs = null;
                AtomicBoolean b = new AtomicBoolean(true);


                if (isRs) {
                    rs = stmt.getResultSet();
                    PrestoResultSet resultSet = rs.unwrap(PrestoResultSet.class);
                    System.out.println("执行 id: " + resultSet.getQueryId());

                    System.out.println(resultSet.getStats().getProgressPercentage());
                    Thread thread = new Thread(() -> {
                        try {
                            while (b.get() && !Thread.interrupted()) {
                                Thread.sleep(10000);
                                System.out.println("是否在排队: " + resultSet.getStats().isQueued());
//                                System.out.println("queryStats: " + queryStatsList.size());
                                System.out.println(resultSet.getStats().getProgressPercentage());
                            }
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                        }

                    });
                    thread.setName("PrestoMonitor-Thread");
                    thread.setDaemon(true);
                    thread.start();

                    while (rs.next()) {
                        b.set(false);
                        System.out.println("拿到结果了............");
                        String res = rs.getString(1);
                        System.out.println(res);
                        // System.out.println(String.format("time=%s", time));
                    }
                }

            } finally {
                stmt.close();
                conn.close();
            }
        } catch (Exception e) {
            e.printStackTrace();
//            log.error("执行presto出错: sql={} 错误信息如下: ", sql, e);
        }


    }

    private static class RecordingProgressMonitor
            implements Consumer<QueryStats> {
        private final ImmutableList.Builder<QueryStats> builder = ImmutableList.builder();
        private boolean finished;

        @Override
        public synchronized void accept(QueryStats queryStats) {
            checkState(!finished);
            System.out.println("monitor状态" + queryStats.getState());
            builder.add(queryStats);
        }

        public synchronized List<QueryStats> finish() {
            finished = true;
            return builder.build();
        }
    }

}
